import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from imageMatching.SuperGlue import SuperGlue
from imageMatching.SuperPoint import SuperPoint
from imageMatching.ImageProcessing import ImageProcessing
import os

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化SuperPoint特征提取模型
superpoint_config = {
    'nms_radius': 4,  # 非极大值抑制半径，用于关键点筛选
    'max_keypoints': 1000,  # 保留的最大关键点数量
    'descriptor_dim': 256,  # 特征描述子的维度
}
superpoint = SuperPoint(superpoint_config).to(device)  # 将模型加载到指定设备

# 初始化SuperGlue特征匹配模型
superglue_config = {
    'weights': 'outdoor',  # 使用室内场景的预训练权重
    'sinkhorn_iterations': 20,  # Sinkhorn算法迭代次数，用于最优匹配计算
    'match_threshold': 0.1,  # 匹配阈值，得分低于此值的匹配将被过滤
}
superglue = SuperGlue(superglue_config).to(device)  # 将模型加载到指定设备

# 定义图像路径：设置可见光和红外图像所在的基础路径
base_path = r"D:\红外线\train"
vis_path = os.path.join(base_path, "vis")  # 可见光图像文件夹
ir_path = os.path.join(base_path, "ir")  # 红外图像文件夹

# 获取所有图像文件：筛选出文件夹中所有常见图像格式的文件
vis_images = [f for f in os.listdir(vis_path) if f.endswith(('.jpg', '.png', '.jpeg'))]
ir_images = [f for f in os.listdir(ir_path) if f.endswith(('.jpg', '.png', '.jpeg'))]

# 确保两个文件夹中的图像数量相同
if len(vis_images) != len(ir_images):
    print(f"警告：可见光图像数量({len(vis_images)})与红外图像数量({len(ir_images)})不匹配")
    min_len = min(len(vis_images), len(ir_images))
    vis_images = vis_images[:min_len]  # 取较小数量的图像作为处理集合
    ir_images = ir_images[:min_len]

# 对每对图像进行处理
for vis_img_name, ir_img_name in zip(vis_images, ir_images):
    print(f"\n处理图像对: {vis_img_name} 和 {ir_img_name}")

    # 读取图像：从指定路径加载可见光和红外图像
    vis_img = cv2.imread(os.path.join(vis_path, vis_img_name))
    ir_img = cv2.imread(os.path.join(ir_path, ir_img_name))

    # 检查图像是否成功加载
    if vis_img is None or ir_img is None:
        print(f"错误：无法读取图像 {vis_img_name} 或 {ir_img_name}")
        continue  # 跳过当前图像对，继续处理下一对

    # 图像预处理：将图像转换为灰度图，调整大小并归一化
    # is_ir=True 表示对红外图像进行对比度增强处理
    vis_padded, vis_info = ImageProcessing.preprocess_image(vis_img, is_ir=False)
    ir_padded, ir_info = ImageProcessing.preprocess_image(ir_img, is_ir=True)

    # 特征提取：使用SuperPoint模型从预处理后的图像中提取关键点和描述子
    vis_kpts, vis_desc = ImageProcessing.extract_keypoints(vis_padded, superpoint, device)
    ir_kpts, ir_desc = ImageProcessing.extract_keypoints(ir_padded, superpoint, device)

    # 准备特征匹配的数据：将关键点和描述子转换为模型所需的张量格式
    data = {
        'keypoints0': torch.from_numpy(vis_kpts).float().to(device).unsqueeze(0),  # 可见光关键点
        'keypoints1': torch.from_numpy(ir_kpts).float().to(device).unsqueeze(0),  # 红外关键点
        'descriptors0': torch.from_numpy(vis_desc).float().to(device).unsqueeze(0).transpose(1, 2),  # 可见光描述子
        'descriptors1': torch.from_numpy(ir_desc).float().to(device).unsqueeze(0).transpose(1, 2),  # 红外描述子
    }

    # 特征匹配：使用SuperGlue模型找出两组关键点之间的匹配关系
    matches = superglue(data)['matches0'][0].cpu().numpy()

    # 筛选有效匹配：过滤掉未匹配的关键点（matches中的-1值表示未匹配）
    valid_mask = (matches != -1)
    matched_vis = vis_kpts[valid_mask]  # 可见光图像中有效的匹配点
    matched_ir = ir_kpts[matches[valid_mask]]  # 红外图像中对应的匹配点

    # 图像配准：使用RANSAC算法计算单应性矩阵，将红外图像对齐到可见光图像
    if len(matched_vis) >= 4:  # 至少需要4对匹配点才能计算单应性矩阵
        h, _ = cv2.findHomography(matched_ir, matched_vis, cv2.RANSAC, 5.0)
        registered_ir = cv2.warpPerspective(ir_img, h, (vis_img.shape[1], vis_img.shape[0]))

        # 创建输出目录：用于保存配准后的图像
        output_dir = os.path.join(base_path, "registered")
        os.makedirs(output_dir, exist_ok=True)  # 目录存在则不创建，避免错误

        # 保存结果：将配准后的红外图像保存到输出目录
        output_name = f"registered_{ir_img_name}"
        cv2.imwrite(os.path.join(output_dir, output_name), registered_ir)
        print(f"成功：已保存配准结果到 {output_name}")
    else:
        print(f"失败：仅找到{len(matched_vis)}对匹配点，不足配准")
